# arq_app.py
import asyncio
import datetime as dt
from typing import Any

from arq.connections import RedisSettings
from arq import create_pool

REDIS = RedisSettings(host="127.0.0.1", port=6379, database=1)

async def long_task(ctx: dict[Any, Any], num:int, sleeper: float):
    print(f"[{dt.datetime.now()}] {num}: Task entered & started...")

    # Mock call
    await asyncio.sleep(1)
    
    # Fail randomly
    # import random
    # if random.random() > 0.5:
    #     raise RuntimeError(f"[{dt.datetime.now()}] {num}: Task randomly failed...")
    
    # Balance task
    await asyncio.sleep(sleeper)

    print(f"[{dt.datetime.now()}] {num}: Task finished...")



def worker_dispatcher(worker_queue_name: str):
    class ArqWorker:
        # Functions to schedule
        functions = [long_task]

        # ARQ settings
        redis_settings = REDIS
        max_jobs = 5
        job_timeout = 9 * 60 * 60 # 9hrs
        keep_result = 9 * 60 * 60 # 9hrs
        health_check_interval = 5
        retry_jobs = False
        allow_abort_jobs = True

        queue_name = worker_queue_name

    return ArqWorker


async def main():
    redis = await create_pool(settings_=REDIS)

    for idx in range(32):
        await redis.enqueue_job("long_task", idx, 6 * 60 * 60, _queue_name=f"queue{idx % 4}")


if __name__ == "__main__":
    asyncio.run(main())

# start_workers.py
import multiprocessing as mp
import logging

from arq import run_worker
from .arq_app import worker_dispatcher

def start_worker(queue_name: str):
    logging.basicConfig(filename='worker.log', level=logging.INFO,
                        format='%(asctime)s [%(process)d] %(levelname)s: %(message)s')
    logging.info(f'Starting worker for {queue_name}')
    run_worker(worker_dispatcher(queue_name))

if __name__ == "__main__":
    for i in range(4):
        mp.Process(target=start_worker, args=(f"queue{i}",)).start()

# cancellation.py
import asyncio

from arq import create_pool
from arq.jobs import Job
from arq.connections import RedisSettings

REDIS = RedisSettings(host="127.0.0.1", port=6379, database=1)

async def job_status(task_id: str, stop: bool = False) -> None:
    job: Job = Job(
        job_id=task_id,
        redis=await create_pool(settings_=REDIS),
    )

    print(f"ID {task_id}: {await job.status()}")

    if stop:
        result: bool = await job.abort()
        print(f"ID {task_id}: Aborted: {result}")


async def main() -> None:
    await job_status(task_id="258a1b3ec3a84913a42d992b970841da", stop=True)


if __name__ == "__main__":
    asyncio.run(main())